import hydra
from omegaconf import DictConfig
from torch.utils.data import DataLoader
from ecgcmr.imaging.img_dataset.ContrastiveImagingDataset import ContrastiveImageDataset
from torch.profiler import profile, ProfilerActivity

@hydra.main(config_path="ecgcmr/conf", config_name="base")
def main(cfg: DictConfig):
    def setup_dataloader(mode='train', batch_size=64, num_workers=4):
        cfg.dataset.paths.data_imaging_test = '/vol/ada_ssd/users/seliv/saved_tensors/imaging_pretrain/updated_imaging_train_data.npy'
        dataset = ContrastiveImageDataset(cfg, mode=mode)
        dataloader = DataLoader(dataset,
                                batch_size=batch_size,
                                shuffle=(mode == 'train'),
                                num_workers=num_workers,
                                pin_memory=True)
        return dataloader

    dataloader = setup_dataloader()

    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memory=True, record_shapes=True) as prof:
        for i, data in enumerate(dataloader):
            if i >= 10: 
                break

    print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

if __name__ == "__main__":
    main()
